import time
import sys
import requests
import re
import json
import multiprocessing

# 绘图接口
# api = 'http://192.168.5.10:9527/api/drawing'
api = 'https://ai.moyi.vip/api/drawing'
# url = 'http://192.168.5.19:7860'
url = 'http://127.0.0.1:7860'

# ChatGLM接口
# chat_api = 'http://192.168.5.19:8000/'
chat_api = 'http://127.0.0.1:8000/'

# 翻译接口
# transformers_api = 'http://192.168.5.19:9527/translater'
transformers_api = 'http://127.0.0.1:9527/translater'

server_id = 0
task_id = 0
parameters = []
img_response = []


def get_info():
    while True:
        # 执行需要定时执行的代码
        print("开始请求绘图服务器信息")
        try:
            response = requests.get(url + "/moyi/info", timeout=5)
            data = response.json()
            # print(data)
            print('当前绘图服务GPU', data[0]['gpu']['name'])
            return data[0]
        except Exception as e:
            print("获取绘图服务状态失败10秒后重拾：", e)
        # 暂停10秒钟
        time.sleep(10)


def join_api(queue, system_info):
    global server_id
    while True:
        # 执行需要定时执行的代码
        print("定时器已触发！")
        try:
            response = requests.post(api + "/join", json=system_info, timeout=30)
            data = response.json()
            if data['code'] == 1 and data['data']['server_id']:
                server_id = data['data']['server_id']
                # print('服务器注册成功', server_id)
                get_task(queue, data['data'])
            else:
                print('服务器注册失败', data['msg'])
                sys.exit()
        except Exception as e:
            print("加入服务器失败30秒后重拾：", e)
        # 暂停10秒钟
        time.sleep(30)


def get_task(queue, service_data):
    while True:
        # 执行需要定时执行的代码
        print("开始申请任务")
        task = False
        try:
            response = requests.post(api + "/get_task", json=service_data, timeout=30)
            data = response.json()
            # print('data', data)
            # time.sleep(30)
            # continue
            if data['code'] == 1:
                task = start_task(queue, data['data'])
            elif data['code'] == 2:
                print('服务器港：', data['msg'])
            else:
                print('获取任务失败', data['msg'])
                return False
        except Exception as e:
            print("获取任务失败：", e)
        # 暂停10秒钟
        if not task:
            time.sleep(30)


def start_task(queue, data):
    global server_id
    global img_response
    print("开始执行任务")
    try:
        print('data', data)
        headers = {
            'Content-type': 'multipart/form-data;'
        }
        if data['zh_prompt']:
            print('翻译中文提示词')
            param = {'mod': 'zh-en', 'text': data['zh_prompt']}
            response = requests.post(transformers_api, data=param)
            data['prompt'] = data['prompt'] + ',' + response.text
        if data['zh_negative_prompt']:
            print('翻译中文反向词')
            param = {'mod': 'zh-en', 'text': data['zh_negative_prompt']}
            response = requests.post(transformers_api, data=param)
            data['zh_negative_prompt'] = data['zh_negative_prompt'] + ',' + response.text
        param = {
            'sd_model_checkpoint': data['sd_model_checkpoint'],
            'sd_vae': data['sd_vae']
        }
        requests.post(url + "/sdapi/v1/options", json=param)
        param = {
            'save_images': True,
            'height': data['height'],
            'width': data['width'],
            'sampler_name': data['sampler_name'],
            'negative_prompt': data['negative_prompt'],
            'prompt': data['prompt'],
            'cfg_scale': data['cfg_scale'],
            'batch_size': data['batch_size'],
            'steps': data['steps']
        }
        print('开始绘图')
        response = requests.post(url + "/moyi/txt2img", json=param)
        datas = response.json()


        datas['parameters']['server_id'] = server_id
        datas['parameters']['task_id'] = data['id']
        # img_response.append(datas)
        queue.put(datas)
        print('数据已存档，等待异步处理')
        return True
    except Exception as e:
        print("任务执行失败：", e)
        return False


def main(queue):
    while True:
        info = get_info()
        print('已经获取到绘图系统信息开始申请加入服务器')
        join_api(queue, info)
        time.sleep(30)


def chat_coroutine(queue):
    global chat_api
    while True:
        responses = queue.get()
        if responses:
            param = responses['parameters']
            print('开始处理自然语言标题追加故事', param)
            param['story'] = []
            param['title'] = []
            for clip in param['clip']:
                # 匹配clip出括号内的描述和权重
                pattern = r"\((.*?)\)"
                # 去掉括号内的描述和权重，得到自然语言部分
                natural_language = re.sub(pattern, "", clip).strip()
                print('chat_api', chat_api)
                print('natural_language', natural_language)

                response = requests.post(chat_api, headers={
                    'Content-Type': 'application/json'
                }, json={
                    "prompt": "给如下描述编一篇中文故事 \n " + natural_language,
                    "history": []
                })
                print('response', response)

                ret = json.loads(response.text)
                param['story'].append(ret['response'])
                response = requests.post(chat_api, headers={
                    'Content-Type': 'application/json'
                }, json={
                    "prompt": "起一个引人注目的故事标题",
                    "history": ret['history']
                })
                ret = json.loads(response.text)
                title = ret['response']
                title = re.sub('[\"\s]', '', title)  # 去除"和空格
                match = re.search('《(.*)》', title)  # 匹配中文括号内的内容
                if match:
                    title = match.group(1)
                param['title'].append(title)
            response = requests.post(api + "/submit_task", json=param)
            datas = response.json()
            print('任务提交完毕', datas['msg'])


def consumer(queue):
    """消费者进程"""
    while True:
        print(f"消费者进程Consumed")

        time.sleep(2)
        item = queue.get()
        if item:
            print(f"开始消费 {item}")
        if item is None:
            break  # 生产结束，退出循环


def status_work(status_queue):
    while True:
        time.sleep(2)
        item = queue.get()
        if item['status']:
            print(f"当前状态{item['status'], item['task_id'], item['server_id']}")
        else:
            print("无状态>>", item)


queue = multiprocessing.Queue()
status_queue = multiprocessing.Queue()
if __name__ == '__main__':
    p1 = multiprocessing.Process(target=main, args=(queue,))
    p2 = multiprocessing.Process(target=chat_coroutine, args=(queue,))

    # 启动进程
    p1.start()
    p2.start()

    # 等待进程结束
    p1.join()
    p2.join()

    print("All tasks completed")
# 创建队列
